package com.hyw.encryption.advice;

import java.util.Base64;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import com.hyw.encryption.annotation.Encrypted;
import com.hyw.encryption.configuration.EncryptionSkipperConfig;
import com.hyw.encryption.provider.IEncryptionProvider;
import com.hyw.encryption.provider.extractor.IEncryptKeyExtract;
import com.hyw.encryption.utils.JsonUtils;

/**
 * 用于对JSON响应体在输出之前做加密
 * 
 * @author Hongyu
 */
@ControllerAdvice
public class EncryptResponseBodyAdvice implements ResponseBodyAdvice<Object> {

  /**
   * 之所以会出现多种处理器,主要原因是系统在业务上, 可能存在多种不同的算法加解密方式.
   * <p>
   * <b>由于处理顺序的问题,系统会遍历的方式执行</b>
   */
  private Map<String, List<IEncryptionProvider>> encryptors;

  @SuppressWarnings("rawtypes")
  private IEncryptKeyExtract keyExtractor;

  private EncryptionSkipperConfig skipper;

  public EncryptResponseBodyAdvice(List<IEncryptionProvider> encryptors,
      @SuppressWarnings("rawtypes") IEncryptKeyExtract encryptKeyExtract, EncryptionSkipperConfig skipper) {
    Assert.isTrue(!CollectionUtils.isEmpty(encryptors), "至少要有一个加密处理器");
    this.encryptors = encryptors.stream().collect(Collectors.groupingBy(IEncryptionProvider::getAlgorithm));
    this.keyExtractor = encryptKeyExtract;
    this.skipper = skipper;
  }

  @Override
  public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
    return returnType.hasMethodAnnotation(Encrypted.class);
  }

  @SuppressWarnings("unchecked")
  @Override
  public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType,
      Class<? extends HttpMessageConverter<?>> selectedConverterType, ServerHttpRequest request,
      ServerHttpResponse response) {
    Encrypted annotaion = returnType.getMethodAnnotation(Encrypted.class);
    Assert.isTrue(StringUtils.hasText(annotaion.value()), "Encrypted注解中的算法值不允许为空");
    Assert.isTrue(encryptors.containsKey(annotaion.value()), String.format("找不到算法类型为[%s]对应的加密处理器", annotaion.value()));
    Objects.requireNonNull(body, "对于加密的响应体不允许为空");

    if (skipper.getEnable() && Objects.equals(request.getHeaders().getFirst(skipper.getKey()), skipper.getValue())) {
      return body;
    } else {
      byte[] original = JsonUtils.toJsonBytes(body);
      byte[] data = original;
      String key = null;
      if (StringUtils.hasText(annotaion.key())) {
        key = keyExtractor.extract(request.getHeaders(), annotaion.key());
        Assert.isTrue(StringUtils.hasText(key), String.format("请求头上缺少对应的参数[%s]", annotaion.key()));
      }
      for (IEncryptionProvider processor : encryptors.get(annotaion.value())) {
        data = processor.encryptByKey(key, data);
      }
      return Base64.getEncoder().encodeToString(data);
    }
  }
}
